package com.luobd.db.impl;

import com.luobd.db.Column;
import com.luobd.db.MetadataFetcher;
import com.luobd.db.Table;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

public class MySQLMetadataFetcher implements MetadataFetcher {


    private final Connection connection;
    private final String catalog;

    public MySQLMetadataFetcher(Connection connection, String catalog) {
        this.connection = connection;
        this.catalog = catalog;
    }

    @Override
    public List<Table> getTables() throws SQLException {
        DatabaseMetaData metaData = connection.getMetaData();
        List<Table> tables = new ArrayList<>();

        try (ResultSet tablesRs = metaData.getTables(catalog, null, "%", new String[]{"TABLE"})) {
            while (tablesRs.next()) {
                Table table = new Table();
                table.setName(tablesRs.getString("TABLE_NAME"));
                table.setComment(getTableCommentDirectly(table.getName()));

                // 获取主键
                List<String> primaryKeys = getPrimaryKeys(table.getName());
                table.setPrimaryKeys(primaryKeys);

                // 获取列信息
                List<Column> columns = getColumns(table.getName(), primaryKeys);
                table.setColumns(columns);

                tables.add(table);
            }
        }
        return tables;
    }

    private List<String> getPrimaryKeys(String tableName) throws SQLException {
        List<String> keys = new ArrayList<>();
        DatabaseMetaData metaData = connection.getMetaData();
        try (ResultSet rs = metaData.getPrimaryKeys(catalog, null, tableName)) {
            while (rs.next()) {
                keys.add(rs.getString("COLUMN_NAME"));
            }
        }
        return keys;
    }


    private String getTableCommentDirectly(String tableName) throws SQLException {
        String sql = "SELECT TABLE_COMMENT " +
                "FROM information_schema.TABLES " +
                "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?";
        try (PreparedStatement ps = connection.prepareStatement(sql)) {
            ps.setString(1, catalog);
            ps.setString(2, tableName);
            try (ResultSet rs = ps.executeQuery()) {
                return rs.next() ? rs.getString("TABLE_COMMENT") : "";
            }
        }
    }

    private List<Column> getColumns(String tableName, List<String> primaryKeys) throws SQLException {
        List<Column> columns = new ArrayList<>();
        DatabaseMetaData metaData = connection.getMetaData();
        try (ResultSet rs = metaData.getColumns(catalog, null, tableName, "%")) {
            while (rs.next()) {
                Column col = new Column();
                col.setName(rs.getString("COLUMN_NAME"));
                col.setType(rs.getString("TYPE_NAME"));
                col.setSize(rs.getInt("COLUMN_SIZE"));
                col.setNullable(rs.getInt("NULLABLE") == DatabaseMetaData.columnNullable);
                col.setComment(rs.getString("REMARKS"));
                col.setPrimaryKey(primaryKeys.contains(col.getName()));
                columns.add(col);
            }
        }
        return columns;
    }

}
